from langchain.embeddings.base import Embeddings
from openai import OpenAI
import os

class DashScopeEmbeddings(Embeddings):
    def __init__(self, api_key, model="text-embedding-v4", dimensions=1024):
        """
        初始化DashScope嵌入模型
        :param api_key: 访问密钥
        :param model: 模型名称
        :param dimensions: 嵌入向量维度
        """
        self.api_key = api_key
        self.model = model
        self.dimensions = dimensions
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

    def embed_documents(self, texts):
        """为文档列表生成嵌入向量（批量处理）"""
        try:
            batch_size = 10  # 控制批量大小
            all_embeddings = []

            for i in range(0, len(texts), batch_size):
                batch = texts[i:i + batch_size]
                completion = self.client.embeddings.create(
                    model=self.model,
                    input=batch,
                    dimensions=self.dimensions,
                    encoding_format="float"
                )
                all_embeddings.extend([data.embedding for data in completion.data])

            return all_embeddings
        except Exception as e:
            print(f"嵌入生成出错: {e}")
            return []

    def embed_query(self, text):
        """为单个查询生成嵌入向量"""
        return self.embed_documents([text])[0]